#include "mem2reg.h"
#include <tuple>

static std::set<Instruction> needEraseInsSet;//需要删除的指令集合
static std::map<int, int> mem2reg_map;    //%r2 = load i32 %r0，记录（r2,r0的上一个store的）<旧的寄存器号，新的寄存器号>
static std::set<int> need_phi_allocs;//普通的alloc，需要插入Phi  
static std::map<PhiInstruction *, int> phi_map;//记录Phi函数指令与对应alloc的寄存器号

// 检查该条alloca指令是否可以被mem2reg
// eg. 数组不可以mem2reg
// eg. 如果该指针直接被使用不可以mem2reg(在SysY一般不可能发生,SysY不支持指针语法)
bool Mem2RegPass::IsPromotable(CFG *C, Instruction AllocaInst) { 
    //TODO("IsPromotable"); 
    // 首先判断是否是alloca指令
    if (AllocaInst->GetOpcode()!= BasicInstruction::ALLOCA) {
        return false;
    }
    auto allocaIns=(AllocaInstruction *)AllocaInst;
    //数组
    if (!(allocaIns->GetDims().empty())) {
            return false;
    } 
    return true;   
}
/*
    int a1 = 5,a2 = 3,a3 = 11,b = 4
    return b // a1,a2,a3 is useless
-----------------------------------------------
pseudo IR is:
    %r0 = alloca i32 ;a1
    %r1 = alloca i32 ;a2
    %r2 = alloca i32 ;a3
    %r3 = alloca i32 ;b
    store 5 -> %r0 ;a1 = 5
    store 3 -> %r1 ;a2 = 3
    store 11 -> %r2 ;a3 = 11
    store 4 -> %r3 ;b = 4
    %r4 = load i32 %r3
    ret i32 %r4
--------------------------------------------------
%r0,%r1,%r2只有store, 但没有load,所以可以删去
优化后的IR(pseudo)为:
    %r3 = alloca i32
    store 4 -> %r3
    %r4 - load i32 %r3
    ret i32 %r4
*/

// vset存储的应该是那些只进行存储但未被加载的内存分配寄存器编号
// 该函数对你的时间复杂度有一定要求, 你需要保证你的时间复杂度小于等于O(nlognlogn), n为该函数的指令数
// 提示:deque直接在中间删除是O(n)的, 可以先标记要删除的指令, 最后想一个快速的方法统一删除
//移除没有 users 的 alloca
void Mem2RegPass::Mem2RegNoUseAlloca(CFG *C, std::set<int> &vset) {//在这删store
    // this function is used in InsertPhi
    //TODO("Mem2RegNoUseAlloca");
    //遍历所有指令，找到一个store指令就看它的目标寄存器号在不在vset里面，如果在这条store指令就要删除
    for (auto [bid, block] : *C->block_map) {
        for (auto Ins : block->Instruction_list) {
            if (Ins->GetOpcode() == BasicInstruction::STORE) {//store <ty> <value>, ptr<pointer>|||store i32 4,ptr %3
                auto storeIns = (StoreInstruction *)Ins;
                if (storeIns->GetPointer()->GetOperandType() == BasicOperand::REG &&
                    vset.find(storeIns->GetPointerRegNo())!= vset.end()) {
                    needEraseInsSet.insert(Ins);
                }
            }
        }
    }

}

/*
    int b = getint();
    b = b + 10
    return b // def and use of b are in same block
-----------------------------------------------
pseudo IR is:
    %r0 = alloca i32 ;b
    %r1 = call getint()
    store %r1 -> %r0
    %r2 = load i32 %r0
    %r3 = %r2 + 10
    store %r3 -> %r0
    %r4 = load i32 %r0
    ret i32 %r4
--------------------------------------------------
%r0的所有load和store都在同一个基本块内
优化后的IR(pseudo)为:
    %r1 = call getint()
    %r3 = %r1 + 10
    ret %r3

对于每一个load，我们只需要找到最近的store,然后用store的值替换之后load的结果即可
*/

// vset是load and store在同一个块里面的alloca寄存器号，如0
// 该函数对你的时间复杂度有一定要求，你需要保证你的时间复杂度小于等于O(nlognlogn), n为该函数的指令数
//如果某个 alloca 的所有 defs 和 users 都在同一个基本块内，且 store 在 load 前面，则可以将基本块内的 load 替换成对应的 store
void Mem2RegPass::Mem2RegUseDefInSameBlock(CFG *C, std::set<int> &vset, int block_id) {
    // this function is used in InsertPhi
    //TODO("Mem2RegUseDefInSameBlock");
    std::map<int, int> alloc_store_map;    //存(alloc reg，现在store的reg)，如（r0,r1)
    //遍历块里面的所有指令
    for (auto Ins : (*C->block_map)[block_id]->Instruction_list) {
        //如果store指令，store %r1 -> %r0，这个store指令就可以删了，但是需要记录一下%r1(都记录了，r0r1键值对)
        if (Ins->GetOpcode() == BasicInstruction::STORE) {
            auto storeIns = (StoreInstruction *)Ins;
            if (storeIns->GetPointer()->GetOperandType() == BasicOperand::REG &&
                vset.find(storeIns->GetPointerRegNo())!= vset.end()) {
                alloc_store_map[storeIns->GetPointerRegNo()] = storeIns->GetValueRegNo();
                needEraseInsSet.insert(Ins);
            }
        }
        //如果是load指令，如%r2 = load i32 %r0，这个load指令就可以删了，但是需要记录一下%r2      
        if (Ins->GetOpcode() == BasicInstruction::LOAD) {
            auto loadIns = (LoadInstruction *)Ins;
            if (loadIns->GetPointer()->GetOperandType() == BasicOperand::REG &&
                vset.find(loadIns->GetPointerRegNo())!= vset.end()) {
                mem2reg_map[loadIns->GetResultRegNo()] = alloc_store_map[loadIns->GetPointerRegNo()];
                needEraseInsSet.insert(Ins);
            }
        }
    }
}

// vset 只有一个定义store支配所有使用load的alloc reg
// 该函数对你的时间复杂度有一定要求，你需要保证你的时间复杂度小于等于O(nlognlogn)
//如果一个 alloca 只有一个 def，那么 users 可以替换成这个 store 的值。
//满足：如果 load 和 store 在同一个基本块，则 store 应该在 load 前面；如果二者在不同基本块，则需要保证 load 指令能被 store 支配
void Mem2RegPass::Mem2RegOneDefDomAllUses(CFG *C, std::set<int> &vset) {
    // this function is used in InsertPhi
    //TODO("Mem2RegOneDefDomAllUses");
    std::map<int, int> alloc_store_value_map;//(alloc,store_value)
    //遍历cfg中所有块的所有指令
    for (auto [bid, block] : *C->block_map) {
        for (auto Ins : block->Instruction_list) {
            if (Ins->GetOpcode() == BasicInstruction::STORE &&
                ((StoreInstruction*)Ins)->GetPointer()->GetOperandType() == BasicOperand::REG &&
                vset.find(((StoreInstruction*)Ins)->GetPointerRegNo())!= vset.end()) {
                alloc_store_value_map[((StoreInstruction*)Ins)->GetPointerRegNo()] = ((StoreInstruction*)Ins)->GetValueRegNo();
                needEraseInsSet.insert(Ins);
            }
        }
    }
    for (auto [bid, block] : *C->block_map) {
        for (auto Ins : block->Instruction_list) {
            if (Ins->GetOpcode() == BasicInstruction::LOAD &&
                     ((LoadInstruction*)Ins)->GetPointer()->GetOperandType() == BasicOperand::REG &&
                     vset.find(((LoadInstruction*)Ins)->GetPointerRegNo())!= vset.end()) {
                mem2reg_map[((LoadInstruction*)Ins)->GetResultRegNo()] = alloc_store_value_map[((LoadInstruction*)Ins)->GetPointerRegNo()];
                needEraseInsSet.insert(Ins);
            }
        }
    }

}

//计算cfg里面所有alloc的定义store和使用load的基本块的编号
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line179-line193
auto CalDefAndUse(CFG *C) {
    //alloc的id，所有定义/使用的位置的基本块
    std::map<int, std::set<int>> defs, uses;
    //alloc的id，数量
    std::map<int, int> def_num;//因为可能一个块里有多个store，所以额外记录一下store的数量
    //遍历cfg所有block的所有指令
    for (auto [bid, block] : *C->block_map) {
        for (auto Ins : block->Instruction_list) {
            if (Ins->GetOpcode() == BasicInstruction::STORE) {
                auto storeIns = (StoreInstruction *)Ins;
                if (storeIns->GetPointer()->GetOperandType()!= BasicOperand::GLOBAL) {
                    defs[storeIns->GetPointerRegNo()].insert(bid);
                    def_num[storeIns->GetPointerRegNo()]++;
                }
            } 
            else if (Ins->GetOpcode() == BasicInstruction::LOAD) {
                auto loadIns = (LoadInstruction *)Ins;
                if (loadIns->GetPointer()->GetOperandType()!= BasicOperand::GLOBAL) {
                    uses[loadIns->GetPointerRegNo()].insert(bid);
                }
            }
        }
    }
    return std::tuple(defs, uses, def_num);
}

void Mem2RegPass::printInsertPhi(CFG *C){
    LLVMBlock entry_Block = (*C->block_map)[0];
     for (auto Ins : entry_Block->Instruction_list) {
        auto allocIns = (AllocaInstruction *)Ins;
        int alloc_regno = allocIns->GetResultRegNo();//获得alloc的寄存器号
std::cout<<"all alloces' regno:"<<alloc_regno<<std::endl;
     }
}

void Mem2RegPass::InsertPhi(CFG *C) { 
    //TODO("InsertPhi"); 
    //printInsertPhi(C);
    auto [defs, uses, def_num] = CalDefAndUse(C);//得到这个cfg中所有store和load的位置
    LLVMBlock entry_Block = (*C->block_map)[0];
    //这里存的就是alloc的id
    std::set<int> no_use_vset, onedom_vset;//没有使用；只有一个定义
    std::map<int, std::set<int>> sameblock_vset_map;//定义和使用在同一个基本块

    //因为把所有的alloc的都插在了函数入口块零号块
    for (auto Ins : entry_Block->Instruction_list) {
        //遍历所有的alloc
        if (Ins->GetOpcode() != BasicInstruction::ALLOCA) {
            continue;
        }
        auto allocIns = (AllocaInstruction *)Ins;
        //先判断这个alloc能不能mem2reg
        if(!IsPromotable(C,allocIns)){
            continue;//不能转换，一般为数组
        }

        int alloc_regno = allocIns->GetResultRegNo();//获得alloc的寄存器号
        //std::cout<<"all alloces' regno:"<<alloc_regno<<std::endl;
        BasicInstruction::LLVMType type = allocIns->GetDataType();

        auto alloc_defs = defs[alloc_regno];//获得这个alloc的所有store的块号，set
        auto alloc_uses = uses[alloc_regno];

        //先对一些alloc进行优化：删掉没有使用的，替换定义和使用在同一个块里面的，只有一个定义支配所有使用的
        if (alloc_uses.size() == 0) {    //没有使用，全都可以删掉
            needEraseInsSet.insert(Ins);
            no_use_vset.insert(alloc_regno);
            continue;
        }

        if (alloc_defs.size() == 1) {//只在一个块里面有定义，可能有多个定义，后续所有的使用都可以替换成那个store的reg
            int block_id = *(alloc_defs.begin());
            if (alloc_uses.size() == 1 && *(alloc_uses.begin()) == block_id) {    //定义和使用在同一个基本块里
                needEraseInsSet.insert(Ins);
                sameblock_vset_map[block_id].insert(alloc_regno);
                continue;
            }
        }

        if (def_num[alloc_regno] == 1) {    //只有一个定义，一个定义支配所有使用
            int block_id = *(alloc_defs.begin());
            int dom_flag = 1;
            //遍历所有使用所在的块
            for (auto load_bid : alloc_uses) {
                if (domtrees->GetDomTree(C)->IsDominate(block_id, load_bid) == false) {//判断该使用能不能被支配
                    dom_flag = 0;
                    break;
                }
            }
            if (dom_flag) {    // 一个定义支配所有使用
                needEraseInsSet.insert(Ins);
                onedom_vset.insert(alloc_regno);
                continue;
            }
        }
        
        // 执行到这里，说明需要插入Phi函数进行处理
        need_phi_allocs.insert(alloc_regno);
        needEraseInsSet.insert(Ins);
        std::set<int> phi_block{};    //用于记录已经插入了Phi函数的基本块编号，避免在同一个基本块重复插入一个alloc的Phi
        std::set<int> alloc_defs2 = defs[alloc_regno];    //alloc_regno的定义所在基本块号们（store)

//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line179-line193
        while (!alloc_defs2.empty()) {
            int bid1 = *alloc_defs2.begin();//定义所在的块们（store)set有顺序吗？这个begin是第一次store?
            alloc_defs2.erase(bid1);
            //遍历基本块bid1的支配边界
            for (auto bid2 : domtrees->GetDomTree(C)->GetDF(bid1)) {
                //std::cout<<alloc_regno<<" "<<bid1<<" "<<bid2<<"\n";
                if (phi_block.find(bid2) == phi_block.end()) {//如果还没有被插入过Phi
                //std::cout<<"max_reg:"<<C->max_reg<<std::endl;
                    PhiInstruction *phiIns = new PhiInstruction(type, GetNewRegOperand(++C->max_reg));
                    (*C->block_map)[bid2]->InsertInstruction(0, phiIns);//在这个块开头处插入Phi指令
                    phi_map[phiIns] = alloc_regno;
                    phi_block.insert(bid2);//这个快已经插入过phi指令
                    // 后续还需要继续检查这个基本块相关的支配边界情况，继续传播和插入Phi函数等操作，直到满足条件不再需要插入为止。
                    if (defs[alloc_regno].find(bid2) == defs[alloc_regno].end()) {
                        alloc_defs2.insert(bid2);
                    }
                }
            }
        } 
    }
    Mem2RegNoUseAlloca(C, no_use_vset);
    Mem2RegOneDefDomAllUses(C, onedom_vset);
    for (auto [id, vset] : sameblock_vset_map) {
        Mem2RegUseDefInSameBlock(C, vset, id);
    }
    
}

// 判断给定的指令是否涉及到特定集合S中的寄存器
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line203-line225
int in_allocs(std::set<int> &S, Instruction Ins) {
    int pointer = -1;
    if (Ins->GetOpcode() == BasicInstruction::LOAD) {
        auto loadIns = (LoadInstruction*)(Ins);
        if (loadIns->GetPointer()->GetOperandType() == BasicOperand::REG) {
            pointer = loadIns->GetPointerRegNo();
        }
    } else if (Ins->GetOpcode() == BasicInstruction::STORE) {
        auto storeIns = (StoreInstruction*)(Ins);
        if (storeIns->GetPointer()->GetOperandType() == BasicOperand::REG) {
            pointer = storeIns->GetPointerRegNo();
        }
    }
    return (S.find(pointer)!= S.end())? pointer : -1;
}

void printIncomingVals(const std::map<int, int>& incomingVals) {
    if (incomingVals.empty()) {
        std::cout << "IncomingVals is empty." << std::endl;
        return;
    }

    for (const auto& [key, value] : incomingVals) { // 使用 C++17 的结构化绑定
        std::cout << "Key: " << key << ", Value: " << value << std::endl;
    }
}

//上面只是删了alloc和哪三种特殊情况的load和store，下面要删需要插入phi指令的alloc的store和load
void Mem2RegPass::VarRename(CFG *C) { 
    //TODO("VarRename"); 
    std::map<int, std::map<int, int>> WorkList;    //< blockid, <alloc_reg,val_reg> >
    WorkList.insert({0, std::map<int, int>{}});//从零号块开始
    std::vector<int> block_visit;//用于标记各个基本块是否已经被访问过
    block_visit.resize(C->max_label + 1);

    while (!WorkList.empty()) {
        int work_block = (*WorkList.begin()).first;
        auto IncomingVals = (*WorkList.begin()).second;//接下来要操作的一对替换寄存器号
        //std::cout<<"while begin:";
        //printIncomingVals(IncomingVals);
        WorkList.erase(work_block);
        if (block_visit[work_block]) {
            continue;
        }
        block_visit[work_block] = 1;
        //遍历当前块的所有指令
        //对指令进行一些操作：alloca，直接删除；load，读取 IncomingVals[]；store 或φ，写入 IncomingVals[]
        //在上面插入phi指令的时候就把alloc放到删除集里面了，这时候就可以删除loadstore了，需要记录一下load的result，store的value，好用来替换，mem2reg_map
        for (auto &Ins : (*C->block_map)[work_block]->Instruction_list) {
            int opcode = Ins->GetOpcode();
            if (opcode == BasicInstruction::LOAD || opcode == BasicInstruction::STORE) {
                auto ptrOpType = (opcode == BasicInstruction::LOAD)
                   ? static_cast<LoadInstruction*>(Ins)->GetPointer()->GetOperandType()
                    : static_cast<StoreInstruction*>(Ins)->GetPointer()->GetOperandType();
                if (ptrOpType == BasicOperand::REG) {
                    int v = in_allocs(need_phi_allocs, Ins);
                    if (v >= 0) {
                        if (opcode == BasicInstruction::LOAD) {
                            auto loadIns = static_cast<LoadInstruction*>(Ins);
                            needEraseInsSet.insert(loadIns);
                            mem2reg_map[loadIns->GetResultRegNo()] = IncomingVals[v];
                        } else {
                            auto storeIns = static_cast<StoreInstruction*>(Ins);
                            needEraseInsSet.insert(storeIns);
                            IncomingVals[v] = static_cast<RegOperand*>(storeIns->GetValue())->GetRegNo();
                        }
                    }
                }
            } else if (opcode == BasicInstruction::PHI) {
                auto phiIns = static_cast<PhiInstruction*>(Ins);
                if (needEraseInsSet.find(phiIns) == needEraseInsSet.end()) {
                    auto it = phi_map.find(phiIns);
                    if (it!= phi_map.end()) {
                        IncomingVals[it->second] = phiIns->GetResultRegNo();
                    }
                }
            }
        }

        //printIncomingVals(IncomingVals);

        //遍历完一个基本块后，遍历它的后继，向所有后继的 𝜑中插入当前基本块流向该后继的边
        for (auto hj_block : C->G[work_block]) {
            int bid = hj_block->block_id;
            WorkList.insert({bid, IncomingVals});
            //遍历这个后继块的所有指令
            for (auto Ins : (*C->block_map)[bid]->Instruction_list) {
                if (Ins->GetOpcode() != BasicInstruction::PHI) {
                    break;
                }
                auto phiIns = (PhiInstruction *)Ins;
                // 找到 phi 对应的 alloca
                auto it = phi_map.find(phiIns);
                if (it != phi_map.end()) {
                    int alloc_reg = it->second;
                    if (IncomingVals.find(alloc_reg) == IncomingVals.end()) {
                        needEraseInsSet.insert(phiIns);//即判断是否在前驱基本块中已经处理过该寄存器的替换情况
                        continue;
                    }
                    // 为 phi 添加前驱块到当前块的边
                    
                    //std::cout<<"add phi_edge:"<<alloc_reg<<"->"<<IncomingVals[alloc_reg]<<","<<GetNewLabelOperand(work_block)<<std::endl;
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line289                 
                    phiIns->InsertPhi(GetNewRegOperand(IncomingVals[alloc_reg]), GetNewLabelOperand(work_block));//就是操作那个phi_list
                }
            }
        }
        
    }
    //到这里所有需要删除的指令都处理完了
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line295-line309
    //mem2reg_map存的是load的result，对应的最新的store
    for (auto [id, bb] : *C->block_map) {
        for (auto Ins : bb->Instruction_list) {
            if (Ins->GetOpcode() == BasicInstruction::LOAD && ((LoadInstruction *)Ins)->GetPointer()->GetOperandType() == BasicOperand::REG) {
                auto loadIns = (LoadInstruction *)Ins;
                int result = loadIns->GetResultRegNo();
                if (mem2reg_map.find(result)!= mem2reg_map.end()) {
                    int result2 = mem2reg_map[result];
                    while (mem2reg_map.find(result2)!= mem2reg_map.end()) {//循环替换，直到找到最根本的
                        mem2reg_map[result] = mem2reg_map[result2];
                        result2 = mem2reg_map[result];
                    }
                }
            }
        }
    }

    //处理需要删除的指令
    for (auto [bid, block] : *C->block_map) {
        auto tmp_Instruction_list = block->Instruction_list;
        block->Instruction_list.clear();
        //遍历这个块的全部指令
        for (auto Ins : tmp_Instruction_list) {
            if (needEraseInsSet.find(Ins) == needEraseInsSet.end()) {
                block->InsertInstruction(1, Ins);//如果这个指令在needEraseInsSet，就是需要删除的
            }
            
        }
    }
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line321-line325
    //来替换指令中的寄存器操作数，确保指令中使用的寄存器是经过变量重命名和替换后的正确寄存器
    for (auto B1 : *C->block_map) {
        for (auto Ins : B1.second->Instruction_list) {  
            Ins->ReplaceRegByMap(mem2reg_map);
        }
    }

    needEraseInsSet.clear();
    mem2reg_map.clear();
    need_phi_allocs.clear();
    phi_map.clear();

}

//就是store指令可能store立即数，为了简化，把这些立即数啥的都转成regOprandoptimize\ssa\mem2reg.cc
//Reference: https://github.com/yuhuifishash/SysY/blob/master/optimize/ssa/mem2reg.cc line335-line358
void Mem2RegInit(CFG *C) {
    for (auto& [bid, block] : *C->block_map) {
        // 直接在原指令列表上进行操作，避免额外的临时列表和复制操作
        for (auto it = block->Instruction_list.begin(); it!= block->Instruction_list.end(); ++it) {
            if ((*it)->GetOpcode() == BasicInstruction::STORE) {
                auto storeIns = static_cast<StoreInstruction*>(*it);
                auto val = storeIns->GetValue();
                if (val->GetOperandType() == BasicOperand::IMMI32) {
                    // 原地构造新的算术指令并插入到当前位置之前
                    it = block->Instruction_list.insert(it, new ArithmeticInstruction(BasicInstruction::ADD, BasicInstruction::I32, val, new ImmI32Operand(0), GetNewRegOperand(++C->max_reg)));
                    // 更新 store 指令的操作数
                    storeIns->SetValue(GetNewRegOperand(C->max_reg));
                } else if (val->GetOperandType() == BasicOperand::IMMF32) {
                    it = block->Instruction_list.insert(it, new ArithmeticInstruction(BasicInstruction::FADD, BasicInstruction::FLOAT32, val, new ImmF32Operand(0), GetNewRegOperand(++C->max_reg)));
                    storeIns->SetValue(GetNewRegOperand(C->max_reg));
                }
            }
        }
    }
}



void Mem2RegPass::Mem2Reg(CFG *C) {
    Mem2RegInit(C);//把store的值都改成寄存器变量
    InsertPhi(C);
    VarRename(C);
}

void Mem2RegPass::Execute() {
    for (auto [defI, cfg] : llvmIR->llvm_cfg) {
        Mem2Reg(cfg);
    }
}
